# coding: utf-8
import re
from pathlib import Path
from subprocess import Popen, PIPE, STDOUT, run
from typing import Union, Tuple

import m3u8
import httpx
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad


class DecodeByte:
    # 解密
    @staticmethod
    def do_decode(key: Union[str, bytes], iv: Union[str, bytes], data: Union[str, bytes], method="AES-128") -> bytes:
        if isinstance(key, str):
            key = key.encode('utf-8')
        if isinstance(iv, str):
            iv = iv.encode('utf-8')
        if "AES-128" == method:
            aes = AES.new(key, AES.MODE_CBC, iv)
            if data and (len(data) % 16) != 0:
                data = pad(data, 16)
            return aes.decrypt(data)
        else:
            return b''


def redirected_resolution_url(
        url: str,
        headers: dict = None,
        verify_ssl: bool = False,
        timeout: int = 10,
):
    """
    从 m3u8 文件获取重定向的解析 URL.
    :param url: m3u8文件 URL.
    :param headers: 请求头.
    :param verify_ssl: 是否验证 SSL 证书.
    :param timeout: 超时时间.
    :return: 重定向的解析 URL, 获取失败返回原 URL.
    """
    headers = headers or {}
    m3u8_obj = m3u8.load(url, headers=headers, verify_ssl=verify_ssl, timeout=timeout)
    # 找到最大分辨率的播放列表
    playlists = m3u8_obj.playlists
    playlist: m3u8.Playlist
    if not playlists:
        return url
    # 获取最大分辨率的播放列表
    __max_resolution = (0, 0)
    max_index = -1
    for i, playlist in enumerate(playlists):
        resolution = playlist.stream_info.resolution
        if not resolution:
            break
        if resolution > __max_resolution:
            __max_resolution = resolution
            max_index = i
    max_url = m3u8_obj.playlists[max_index].absolute_uri
    return max_url


def m3u8_to_absolute_uris(
        url: str,
        headers: dict = None,
        verify_ssl: bool = False,
        timeout: int = 10
):
    """
    从 m3u8 文件获取重定向的解析 URL 列表.
    :param url: m3u8文件 URL.
    :param headers: 请求头.
    :param verify_ssl: 是否验证SSL证书.
    :param timeout: 超时时间.
    :return: 重定向的解析URL列表.
    """
    headers = headers or {}
    m3u8_obj = m3u8.load(url, headers=headers, verify_ssl=verify_ssl, timeout=timeout)
    absolute_uris = []
    kt: m3u8.Key = m3u8_obj.segments[0].key
    key = {
        'key': '',
        'method': '',
        'uri': '',
        'iv': '',
        'keyformat': '',
        'keyformatversions': ''
    }
    if kt:
        response = httpx.get(kt.absolute_uri, headers=headers, verify=verify_ssl)
        response.raise_for_status()
        iv = kt.iv
        if iv:
            iv = iv[2:18]
        key = {
            'key': response.content,
            'method': kt.method,
            'uri': kt.absolute_uri,
            'iv': iv,
            'keyformat': kt.keyformat,
            'keyformatversions': kt.keyformatversions
        }
    for segment in m3u8_obj.segments:
        segment: m3u8.Segment
        absolute_uris.append(segment.absolute_uri)
    return absolute_uris, key


def sortFfmpegLog(ffmpeg_file: Union[str, Path]):
    """
    排序 ffmpeg 日志信息
    :return:
    """
    files = ffmpeg_file.read_text(encoding='utf-8').split('\n')
    # 去除空行
    files = list(filter(None, files))

    # 按文件名排序
    def file_sort(x):
        return re.findall(r'\d+', x)[0]

    files.sort(key=file_sort)
    ffmpeg_file.write_text('\n'.join(files), encoding='utf-8')


def merge_ts_to_mp4(
        ts_path: Union[str, Path],
        out_path: Union[str, Path],
        ffmpeg: str = 'ffmpeg',
) -> Tuple[bool, str]:
    """
    合并 ts 文件为 mp4 文件.
    :param ts_path: ts 文件路径.
    :param out_path: 输出文件路径.
    :param ffmpeg: ffmpeg 可执行文件路径.
    :return:
    """
    if not isinstance(ts_path, Path):
        ts_path = Path(ts_path)
    if not isinstance(out_path, Path):
        out_path = Path(out_path)

    sortFfmpegLog(ts_path)
    cmd = [ffmpeg, '-f', 'concat', '-safe', '0', '-i', ts_path.as_posix(), '-c', 'copy', out_path.as_posix()]
    try:
        process = Popen(cmd, stdout=PIPE, stderr=PIPE, text=True)
        process.wait()
        if process.returncode != 0:
            return False, str(process.stderr)
        else:
            return True, 'ts 文件合并成功'
    except FileNotFoundError as e:
        return False, str(e)
